// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#if defined(EIGEN_USE_THREADS) && !defined(EIGEN_CXX11_TENSOR_TENSOR_DEVICE_THREAD_POOL_H)
#define EIGEN_CXX11_TENSOR_TENSOR_DEVICE_THREAD_POOL_H

namespace Eigen {

// Runs an arbitrary function and then calls Notify() on the passed in
// Notification.
template<typename Function, typename... Args>
struct FunctionWrapperWithNotification
{
	static void run(Notification* n, Function f, Args... args)
	{
		f(args...);
		if (n) {
			n->Notify();
		}
	}
};

template<typename Function, typename... Args>
struct FunctionWrapperWithBarrier
{
	static void run(Barrier* b, Function f, Args... args)
	{
		f(args...);
		if (b) {
			b->Notify();
		}
	}
};

template<typename SyncType>
static EIGEN_STRONG_INLINE void
wait_until_ready(SyncType* n)
{
	if (n) {
		n->Wait();
	}
}

// An abstract interface to a device specific memory allocator.
class Allocator
{
  public:
	virtual ~Allocator() {}
	virtual void* allocate(size_t num_bytes) const = 0;
	virtual void deallocate(void* buffer) const = 0;
};

// Build a thread pool device on top the an existing pool of threads.
struct ThreadPoolDevice
{
	// The ownership of the thread pool remains with the caller.
	ThreadPoolDevice(ThreadPoolInterface* pool, int num_cores, Allocator* allocator = nullptr)
		: pool_(pool)
		, num_threads_(num_cores)
		, allocator_(allocator)
	{
	}

	EIGEN_STRONG_INLINE void* allocate(size_t num_bytes) const
	{
		return allocator_ ? allocator_->allocate(num_bytes) : internal::aligned_malloc(num_bytes);
	}

	EIGEN_STRONG_INLINE void deallocate(void* buffer) const
	{
		if (allocator_) {
			allocator_->deallocate(buffer);
		} else {
			internal::aligned_free(buffer);
		}
	}

	EIGEN_STRONG_INLINE void* allocate_temp(size_t num_bytes) const { return allocate(num_bytes); }

	EIGEN_STRONG_INLINE void deallocate_temp(void* buffer) const { deallocate(buffer); }

	template<typename Type>
	EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Type get(Type data) const
	{
		return data;
	}

	EIGEN_STRONG_INLINE void memcpy(void* dst, const void* src, size_t n) const
	{
#ifdef __ANDROID__
		::memcpy(dst, src, n);
#else
		// TODO(rmlarsen): Align blocks on cache lines.
		// We have observed that going beyond 4 threads usually just wastes
		// CPU cycles due to the threads competing for memory bandwidth, so we
		// statically schedule at most 4 block copies here.
		const size_t kMinBlockSize = 32768;
		const size_t num_threads = CostModel::numThreads(n, TensorOpCost(1.0, 1.0, 0), 4);
		if (n <= kMinBlockSize || num_threads < 2) {
			::memcpy(dst, src, n);
		} else {
			const char* src_ptr = static_cast<const char*>(src);
			char* dst_ptr = static_cast<char*>(dst);
			const size_t blocksize = (n + (num_threads - 1)) / num_threads;
			Barrier barrier(static_cast<int>(num_threads - 1));
			// Launch the last 3 blocks on worker threads.
			for (size_t i = 1; i < num_threads; ++i) {
				enqueue_with_barrier(&barrier, [n, i, src_ptr, dst_ptr, blocksize] {
					::memcpy(
						dst_ptr + i * blocksize, src_ptr + i * blocksize, numext::mini(blocksize, n - (i * blocksize)));
				});
			}
			// Launch the first block on the main thread.
			::memcpy(dst_ptr, src_ptr, blocksize);
			barrier.Wait();
		}
#endif
	}
	EIGEN_STRONG_INLINE void memcpyHostToDevice(void* dst, const void* src, size_t n) const { memcpy(dst, src, n); }
	EIGEN_STRONG_INLINE void memcpyDeviceToHost(void* dst, const void* src, size_t n) const { memcpy(dst, src, n); }

	EIGEN_STRONG_INLINE void memset(void* buffer, int c, size_t n) const { ::memset(buffer, c, n); }

	EIGEN_STRONG_INLINE int numThreads() const { return num_threads_; }

	// Number of theads available in the underlying thread pool. This number can
	// be different from the value returned by numThreads().
	EIGEN_STRONG_INLINE int numThreadsInPool() const { return pool_->NumThreads(); }

	EIGEN_STRONG_INLINE size_t firstLevelCacheSize() const { return l1CacheSize(); }

	EIGEN_STRONG_INLINE size_t lastLevelCacheSize() const
	{
		// The l3 cache size is shared between all the cores.
		return l3CacheSize() / num_threads_;
	}

	EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE int majorDeviceVersion() const
	{
		// Should return an enum that encodes the ISA supported by the CPU
		return 1;
	}

	template<class Function, class... Args>
	EIGEN_STRONG_INLINE Notification* enqueue(Function&& f, Args&&... args) const
	{
		Notification* n = new Notification();
		pool_->Schedule(std::bind(&FunctionWrapperWithNotification<Function, Args...>::run, n, std::move(f), args...));
		return n;
	}

	template<class Function, class... Args>
	EIGEN_STRONG_INLINE void enqueue_with_barrier(Barrier* b, Function&& f, Args&&... args) const
	{
		pool_->Schedule(std::bind(&FunctionWrapperWithBarrier<Function, Args...>::run, b, std::move(f), args...));
	}

	template<class Function, class... Args>
	EIGEN_STRONG_INLINE void enqueueNoNotification(Function&& f, Args&&... args) const
	{
		if (sizeof...(args) > 0) {
			pool_->Schedule(std::bind(std::move(f), args...));
		} else {
			pool_->Schedule(std::move(f));
		}
	}

	// Returns a logical thread index between 0 and pool_->NumThreads() - 1 if
	// called from one of the threads in pool_. Returns -1 otherwise.
	EIGEN_STRONG_INLINE int currentThreadId() const { return pool_->CurrentThreadId(); }

	// WARNING: This function is synchronous and will block the calling thread.
	//
	// Synchronous parallelFor executes f with [0, n) arguments in parallel and
	// waits for completion. F accepts a half-open interval [first, last). Block
	// size is chosen based on the iteration cost and resulting parallel
	// efficiency. If block_align is not nullptr, it is called to round up the
	// block size.
	void parallelFor(Index n,
					 const TensorOpCost& cost,
					 std::function<Index(Index)> block_align,
					 std::function<void(Index, Index)> f) const
	{
		if (EIGEN_PREDICT_FALSE(n <= 0)) {
			return;
			// Compute small problems directly in the caller thread.
		} else if (n == 1 || numThreads() == 1 || CostModel::numThreads(n, cost, static_cast<int>(numThreads())) == 1) {
			f(0, n);
			return;
		}

		// Compute block size and total count of blocks.
		ParallelForBlock block = CalculateParallelForBlock(n, cost, block_align);

		// Recursively divide size into halves until we reach block_size.
		// Division code rounds mid to block_size, so we are guaranteed to get
		// block_count leaves that do actual computations.
		Barrier barrier(static_cast<unsigned int>(block.count));
		std::function<void(Index, Index)> handleRange;
		handleRange = [=, &handleRange, &barrier, &f](Index firstIdx, Index lastIdx) {
			while (lastIdx - firstIdx > block.size) {
				// Split into halves and schedule the second half on a different thread.
				const Index midIdx = firstIdx + divup((lastIdx - firstIdx) / 2, block.size) * block.size;
				pool_->Schedule([=, &handleRange]() { handleRange(midIdx, lastIdx); });
				lastIdx = midIdx;
			}
			// Single block or less, execute directly.
			f(firstIdx, lastIdx);
			barrier.Notify();
		};

		if (block.count <= numThreads()) {
			// Avoid a thread hop by running the root of the tree and one block on the
			// main thread.
			handleRange(0, n);
		} else {
			// Execute the root in the thread pool to avoid running work on more than
			// numThreads() threads.
			pool_->Schedule([=, &handleRange]() { handleRange(0, n); });
		}

		barrier.Wait();
	}

	// Convenience wrapper for parallelFor that does not align blocks.
	void parallelFor(Index n, const TensorOpCost& cost, std::function<void(Index, Index)> f) const
	{
		parallelFor(n, cost, nullptr, std::move(f));
	}

	// WARNING: This function is asynchronous and will not block the calling thread.
	//
	// Asynchronous parallelFor executes f with [0, n) arguments in parallel
	// without waiting for completion. When the last block finished, it will call
	// 'done' callback. F accepts a half-open interval [first, last). Block size
	// is chosen based on the iteration cost and resulting parallel efficiency. If
	// block_align is not nullptr, it is called to round up the block size.
	void parallelForAsync(Index n,
						  const TensorOpCost& cost,
						  std::function<Index(Index)> block_align,
						  std::function<void(Index, Index)> f,
						  std::function<void()> done) const
	{
		// Compute small problems directly in the caller thread.
		if (n <= 1 || numThreads() == 1 || CostModel::numThreads(n, cost, static_cast<int>(numThreads())) == 1) {
			f(0, n);
			done();
			return;
		}

		// Compute block size and total count of blocks.
		ParallelForBlock block = CalculateParallelForBlock(n, cost, block_align);

		ParallelForAsyncContext* const ctx = new ParallelForAsyncContext(block.count, std::move(f), std::move(done));

		// Recursively divide size into halves until we reach block_size.
		// Division code rounds mid to block_size, so we are guaranteed to get
		// block_count leaves that do actual computations.
		ctx->handle_range = [this, ctx, block](Index firstIdx, Index lastIdx) {
			while (lastIdx - firstIdx > block.size) {
				// Split into halves and schedule the second half on a different thread.
				const Index midIdx = firstIdx + divup((lastIdx - firstIdx) / 2, block.size) * block.size;
				pool_->Schedule([ctx, midIdx, lastIdx]() { ctx->handle_range(midIdx, lastIdx); });
				lastIdx = midIdx;
			}

			// Single block or less, execute directly.
			ctx->f(firstIdx, lastIdx);

			// Delete async context if it was the last block.
			if (ctx->count.fetch_sub(1) == 1)
				delete ctx;
		};

		if (block.count <= numThreads()) {
			// Avoid a thread hop by running the root of the tree and one block on the
			// main thread.
			ctx->handle_range(0, n);
		} else {
			// Execute the root in the thread pool to avoid running work on more than
			// numThreads() threads.
			pool_->Schedule([ctx, n]() { ctx->handle_range(0, n); });
		}
	}

	// Convenience wrapper for parallelForAsync that does not align blocks.
	void parallelForAsync(Index n,
						  const TensorOpCost& cost,
						  std::function<void(Index, Index)> f,
						  std::function<void()> done) const
	{
		parallelForAsync(n, cost, nullptr, std::move(f), std::move(done));
	}

	// Thread pool accessor.
	ThreadPoolInterface* getPool() const { return pool_; }

	// Allocator accessor.
	Allocator* allocator() const { return allocator_; }

  private:
	typedef TensorCostModel<ThreadPoolDevice> CostModel;

	// For parallelForAsync we must keep passed in closures on the heap, and
	// delete them only after `done` callback finished.
	struct ParallelForAsyncContext
	{
		ParallelForAsyncContext(Index block_count,
								std::function<void(Index, Index)> block_f,
								std::function<void()> done_callback)
			: count(block_count)
			, f(std::move(block_f))
			, done(std::move(done_callback))
		{
		}
		~ParallelForAsyncContext() { done(); }

		std::atomic<Index> count;
		std::function<void(Index, Index)> f;
		std::function<void()> done;

		std::function<void(Index, Index)> handle_range;
	};

	struct ParallelForBlock
	{
		Index size;	 // block size
		Index count; // number of blocks
	};

	// Calculates block size based on (1) the iteration cost and (2) parallel
	// efficiency. We want blocks to be not too small to mitigate parallelization
	// overheads; not too large to mitigate tail effect and potential load
	// imbalance and we also want number of blocks to be evenly dividable across
	// threads.
	ParallelForBlock CalculateParallelForBlock(const Index n,
											   const TensorOpCost& cost,
											   std::function<Index(Index)> block_align) const
	{
		const double block_size_f = 1.0 / CostModel::taskSize(1, cost);
		const Index max_oversharding_factor = 4;
		Index block_size =
			numext::mini(n, numext::maxi<Index>(divup<Index>(n, max_oversharding_factor * numThreads()), block_size_f));
		const Index max_block_size = numext::mini(n, 2 * block_size);

		if (block_align) {
			Index new_block_size = block_align(block_size);
			eigen_assert(new_block_size >= block_size);
			block_size = numext::mini(n, new_block_size);
		}

		Index block_count = divup(n, block_size);

		// Calculate parallel efficiency as fraction of total CPU time used for
		// computations:
		double max_efficiency =
			static_cast<double>(block_count) / (divup<int>(block_count, numThreads()) * numThreads());

		// Now try to increase block size up to max_block_size as long as it
		// doesn't decrease parallel efficiency.
		for (Index prev_block_count = block_count; max_efficiency < 1.0 && prev_block_count > 1;) {
			// This is the next block size that divides size into a smaller number
			// of blocks than the current block_size.
			Index coarser_block_size = divup(n, prev_block_count - 1);
			if (block_align) {
				Index new_block_size = block_align(coarser_block_size);
				eigen_assert(new_block_size >= coarser_block_size);
				coarser_block_size = numext::mini(n, new_block_size);
			}
			if (coarser_block_size > max_block_size) {
				break; // Reached max block size. Stop.
			}
			// Recalculate parallel efficiency.
			const Index coarser_block_count = divup(n, coarser_block_size);
			eigen_assert(coarser_block_count < prev_block_count);
			prev_block_count = coarser_block_count;
			const double coarser_efficiency = static_cast<double>(coarser_block_count) /
											  (divup<int>(coarser_block_count, numThreads()) * numThreads());
			if (coarser_efficiency + 0.01 >= max_efficiency) {
				// Taking it.
				block_size = coarser_block_size;
				block_count = coarser_block_count;
				if (max_efficiency < coarser_efficiency) {
					max_efficiency = coarser_efficiency;
				}
			}
		}

		return { block_size, block_count };
	}

	ThreadPoolInterface* pool_;
	int num_threads_;
	Allocator* allocator_;
};

} // end namespace Eigen

#endif // EIGEN_CXX11_TENSOR_TENSOR_DEVICE_THREAD_POOL_H
